"""
==============================
💡 LangGraph 中的 Send 机制详解
==============================

在 LangGraph 中，Send 是一个**“并行派发机制”**。
它的作用是：**让一个节点可以一次性派发出多个子任务（子节点调用），
每个子任务带不同的输入参数，并在所有子任务结束后自动聚合结果。**

可以把它想成：
🧠「一个调度员节点」负责根据输入状态，生成多个“分身任务”。
每个任务去做不同的工作（比如针对不同主题生成笑话）。
等所有任务完成后，LangGraph 会自动把结果合并回来。

👉 用途示例：
- 给多个主题并行生成笑话
- 给多条数据并行调用 LLM
- 并行执行多个子服务，再统一合并结果

这个 Demo 会演示：
1️⃣ 从 START 节点开始
2️⃣ 根据输入的 topics 列表，动态生成多个 Send("generate_joke") 任务
3️⃣ 每个任务并行生成一个笑话
4️⃣ LangGraph 自动把所有笑话汇总进最终状态里
"""

# ==========================
# LangGraph - Send 机制 Demo
# ==========================

import operator
from typing import Annotated, TypedDict

from langgraph.graph import START, END
from langgraph.graph import StateGraph
from langgraph.types import Send


# 定义全局状态类型
class OverallState(TypedDict):
    """
    OverallState 表示图运行时共享的状态。
    - topics: 主题列表（输入）
    - jokes: 收集到的笑话列表（输出）
    """
    topics: list[str]
    # operator.add 表示多个节点返回的结果会自动拼接
    jokes: Annotated[list[str], operator.add]


def dispatch_to_joke_nodes(state: OverallState):
    """
    动态决定要派发哪些子任务。
    每个 Send 表示一个“将被执行的节点调用”，携带不同的输入参数。
    """
    return [Send("generate_joke", {"topic": topic}) for topic in state["topics"]]


# 创建状态图构建器
builder = StateGraph(OverallState)


# 定义节点逻辑
def make_joke(state):
    """
    生成一个关于某个主题的笑话。
    实际使用中这里可以替换为 LLM 调用或外部 API。
    """
    topic = state["topic"]
    return {"jokes": [f"Here's a funny joke about {topic}: {topic} walks into a bar... 🍺"]}


# 注册节点
builder.add_node("generate_joke", make_joke)

# 配置动态边：START -> dispatch_to_joke_nodes -> generate_joke
builder.add_conditional_edges(START, dispatch_to_joke_nodes)

# 每个子节点执行完后流向 END
builder.add_edge("generate_joke", END)

# 编译图
graph = builder.compile()

# 执行
result = graph.invoke({"topics": ["AI", "Python", "LangChain"]})
print("🧠 最终合并结果：")
print(result)
# 输出：
# {
#   'topics': ['AI', 'Python', 'LangChain'],
#   'jokes': [
#       "Here's a funny joke about AI: AI walks into a bar... 🍺",
#       "Here's a funny joke about Python: Python walks into a bar... 🍺",
#       "Here's a funny joke about LangChain: LangChain walks into a bar... 🍺"
#   ]
# }


# 可视化输出
graph_png = graph.get_graph().draw_mermaid_png()
with open("llm_langgraph_7_send机制.png", "wb") as f:
    f.write(graph_png)
